#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
训练回调模块
提供增强的训练进度回调功能
"""

from transformers.trainer_callback import TrainerCallback


class ProgressCallback(TrainerCallback):
    """增强的训练进度回调"""
    def __init__(self, total_epochs):
        self.total_epochs = total_epochs
        
    def on_train_begin(self, args, state, control, **kwargs):
        print("🚀 开始训练...")
        print(f"📈 总步数: {state.max_steps}")
        
    def on_train_end(self, args, state, control, **kwargs):
        print("✅ 训练完成!")
        
    # def on_epoch_begin(self, args, state, control, **kwargs):
    #     current_epoch = int(state.epoch) + 1
    #     print(f"🔁 开始第 {current_epoch}/{self.total_epochs} 轮训练")
        
    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs:
            loss = logs.get('loss', 'N/A')
            learning_rate = logs.get('learning_rate', 'N/A')
            # 处理可能为字符串的值
            if isinstance(loss, str):
                print(f"📊 Step {state.global_step}: loss={loss}, lr={learning_rate}")
            else:
                print(f"📊 Step {state.global_step}: loss={loss:.4f}, lr={learning_rate:.2e}")


if __name__ == "__main__":
    # 可以添加一些测试代码
    pass